{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TAO Data services\n",
    "\n",
    "TAO Data Services include 4 key pipelines:\n",
    "1. offline data augmentation using DALI\n",
    "2. auto labeling using TAO Mask Auto-labeler (MAL)\n",
    "3. annotation conversion service\n",
    "4. data analytics service\n",
    "\n",
    "The offline data augmentation service enables users to enrich their data with GPU-accelerated spatial/color/kernel augmentation routines with more control than random augmentation often used in online augmentation during training.\n",
    "\n",
    "Annotating an image-based dataset can be quite tedious and time-consuming. Labeling an instance mask by drawing a good polygon around an object can take 10 times longer than a bounding box. The Auto-Label service is designed to automatically generate instance segmentation masks given the groundtruth bounding boxes, which will greatly reduce the labeling effort.\n",
    "\n",
    "Annotation conversion service provides an easy way of converting annotation groundtruth between COCO and KITTI formats, which are widely used in TAO models.\n",
    "\n",
    "The Data Analytics service analyzes object-detection annotation files and image files, calculates insights,\n",
    "and generate a graph and summary based on various metrics like object size, object counts, etc.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/common/mal_sample.jpg?raw=true\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Convert KITTI dataset to COCO format\n",
    "* Run auto-labeling to generate pseudo masks for KITTI bounding boxes\n",
    "* Apply data augmentation to the KITTI dataset with bounding boxe refinement\n",
    "* Run data analytics to collect useful statistics on the original and augmented KITTI dataset\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of MAL using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset](#head-2)\n",
    "3. [Convert KITTI dataset to COCO format](#head-3)\n",
    "4. [Generate pseudo-masks with the auto-labeler](#head-4)\n",
    "5. [Apply data augmentation](#head-5)\n",
    "6. [Run data analytics](#head-6)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/mal/`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data_services\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/data_services\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()), \"specs\")\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()), \"data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "         # Mapping the Local project directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         },\n",
    "        \"network\": \"host\"\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in the `nvidia-pyindex` python index. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-pyindex\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download KITTI dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the kitti object detection dataset for this example. To find more details, please visit http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Please download both, the left color images of the object dataset from [here](http://www.cvlibs.net/download.php?file=data_object_image_2.zip) and, the training labels for the object dataset from [here](http://www.cvlibs.net/download.php?file=data_object_label_2.zip), and place the zip files in `$HOST_DATA_DIR`\n",
    "\n",
    "Once unzipped, the dataset should have the following structure\n",
    "* training images in `$HOST_DATA_DIR/training/image_2`\n",
    "* training labels in `$HOST_DATA_DIR/training/label_2`\n",
    "* testing images in `$HOST_DATA_DIR/testing/image_2`\n",
    "\n",
    "You may use this notebook with your own dataset as well. To use this example with your own dataset, please follow the same directory structure as mentioned below."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2. Using the sample KITTI dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the rest of this notebook, we will a sample KITTI dataset which consists of 10 image/label pairs randomly picked from the original KITTI dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp $HOST_DATA_DIR/training/image_2/00001*.png $HOST_DATA_DIR/images\n",
    "!ls -l $HOST_DATA_DIR/images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp $HOST_DATA_DIR/training/label_2/00001*.txt $HOST_DATA_DIR/labels\n",
    "!ls -l $HOST_DATA_DIR/labels/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Convert KITTI data to COCO format <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Most TAO object detection models take KITTI or COCO annotation as input. Here we demonstrate how to easily convert KITTI data to COCO [format](https://cocodataset.org/#format-data).  The converted COCO json file will be used to generate pseudo-masks in the next section.\n",
    "\n",
    "COCO format:\n",
    "```\n",
    "    annotation{\n",
    "    \"id\": int, \n",
    "    \"image_id\": int, \n",
    "    \"category_id\": int, \n",
    "    \"segmentation\": RLE or [polygon], \n",
    "    \"area\": float, \n",
    "    \"bbox\": [x,y,width,height], \n",
    "    \"iscrowd\": 0 or 1,\n",
    "    }\n",
    "\n",
    "    image{\n",
    "    \"id\": int,\n",
    "    \"width\": int,\n",
    "    \"height\": int,\n",
    "    \"file_name\": str,\n",
    "    \"license\": int,\n",
    "    \"flickr_url\": str,\n",
    "    \"coco_url\": str,\n",
    "    \"date_captured\": datetime,\n",
    "    }\n",
    "\n",
    "    categories[{\n",
    "    \"id\": int, \n",
    "    \"name\": str, \n",
    "    \"supercategory\": str,\n",
    "    }]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/convert.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert KITTI to COCO\n",
    "!tao dataset annotations convert -e $SPECS_DIR/convert.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The converted json file is saved at:\n",
    "!ls -l $HOST_RESULTS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Generate pseudo-masks with the auto-labeler <a class=\"anchor\" id=\"head-4\"></a>\n",
    "Here we will use a pretrained MAL model to generate pseudo-masks for the converted KITTI data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/autolabel.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 Download auto-labeler pretrained model from NGC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm -f $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List available models\n",
    "!ngc registry model list nvidia/tao/mask_auto_label:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the model\n",
    "!ngc registry model download-version nvidia/tao/mask_auto_label:trainable_v1.0 --dest $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/mask_auto_label_vtrainable_v1.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 Generate pseudo-labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change `gpus` in autolabel.yaml based on your machine.\")\n",
    "!tao dataset auto_label generate -e $SPECS_DIR/autolabel.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check the pseudo label:\")\n",
    "!ls -l $HOST_RESULTS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's visualize the pseudo-masks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install deps\n",
    "!pip3 install Cython==0.29.36\n",
    "!pip3 install numpy\n",
    "!pip3 install matplotlib\n",
    "!pip3 install pillow\n",
    "!pip3 install pycocotools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "from pycocotools.coco import COCO\n",
    "%matplotlib inline\n",
    "\n",
    "image_dir = os.path.join(os.environ[\"HOST_DATA_DIR\"], 'images')\n",
    "json_path = os.path.join(os.environ[\"HOST_RESULTS_DIR\"], 'data_mal.json')\n",
    "coco_mal = COCO(annotation_file=json_path)\n",
    "\n",
    "for i in coco_mal.getImgIds()[-2:-1]:\n",
    "    img_info = coco_mal.loadImgs(i)[0]\n",
    "    img_file_name = img_info[\"file_name\"]\n",
    "    print(img_file_name)\n",
    "    ann_ids = coco_mal.getAnnIds(imgIds=[i], iscrowd=None)\n",
    "    anns = coco_mal.loadAnns(ann_ids)\n",
    "    # raw image\n",
    "    im = Image.open(os.path.join(image_dir, img_file_name))\n",
    "    # plots\n",
    "    fig = plt.figure(figsize = (10,10))\n",
    "    ax1 = fig.add_subplot(211)\n",
    "    ax1.imshow(np.asarray(im))\n",
    "    ax2 = fig.add_subplot(212)\n",
    "    ax2.imshow(np.asarray(im), aspect='auto')\n",
    "    coco_mal.showAnns(anns, draw_bbox=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Apply data augmentation <a class=\"anchor\" id=\"head-5\"></a>\n",
    "\n",
    "In this section, we run offline augmentation with the KITTI data. During the augmentation process, we can use the pseudo-masks generated from the last step to refine the distorted or rotated bounding boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/augment.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.a Augmentation without bounding box refinement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change `num_gpus` in augment.yaml based on your machine.\")\n",
    "!tao dataset augmentation generate -e $SPECS_DIR/augment.yaml results_dir=$RESULTS_DIR/augmented"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Augmented KITTI data are saved in `result/augmented` dir**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $HOST_RESULTS_DIR/augmented/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.b Augmentation with bounding box refinement"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change `num_gpus` in augment.yaml based on your machine.\")\n",
    "!tao dataset augmentation generate -e $SPECS_DIR/augment.yaml \\\n",
    "    results_dir=$RESULTS_DIR/refined \\\n",
    "    spatial_aug.rotation.refine_box.enabled=True \\\n",
    "    spatial_aug.rotation.refine_box.gt_cache=$RESULTS_DIR/data_mal.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Augmented images and refined labels are saved in `result/refined` dir**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $HOST_RESULTS_DIR/refined/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.c Let's visualize the augmented and refined data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image, ImageDraw\n",
    "\n",
    "def draw_bbox(image_path, label_path):\n",
    "    img = Image.open(image_path)\n",
    "    draw = ImageDraw.Draw(img)\n",
    "    f = open(label_path, 'r')\n",
    "    for line in f:\n",
    "        po = list(map(lambda x:float(x), line.split(' ')[4:8]))\n",
    "        draw.rectangle(po, outline=\"Yellow\")\n",
    "    return img\n",
    "image_dir = os.path.join(os.environ[\"HOST_DATA_DIR\"], 'images')\n",
    "label_dir = os.path.join(os.environ[\"HOST_DATA_DIR\"], 'labels')\n",
    "# original image\n",
    "image_path = os.path.join(image_dir, '000015.png')\n",
    "label_path = os.path.join(label_dir, '000015.txt')\n",
    "# plot \n",
    "im = draw_bbox(image_path, label_path)\n",
    "fig = plt.figure(figsize = (10,10))\n",
    "ax1 = fig.add_subplot(311)\n",
    "ax1.imshow(np.asarray(im))\n",
    "# augmented image\n",
    "image_path = os.path.join(os.environ[\"HOST_RESULTS_DIR\"], 'augmented/images/000015.png')\n",
    "label_path = os.path.join(os.environ[\"HOST_RESULTS_DIR\"], 'augmented/labels/000015.txt')\n",
    "# plot \n",
    "im = draw_bbox(image_path, label_path)\n",
    "fig = plt.figure(figsize = (10,10))\n",
    "ax2 = fig.add_subplot(312)\n",
    "ax2.imshow(np.asarray(im))\n",
    "# original image\n",
    "image_path = os.path.join(os.environ[\"HOST_RESULTS_DIR\"], 'refined/images/000015.png')\n",
    "label_path = os.path.join(os.environ[\"HOST_RESULTS_DIR\"], 'refined/labels/000015.txt')\n",
    "# plot \n",
    "im = draw_bbox(image_path, label_path)\n",
    "fig = plt.figure(figsize = (10,10))\n",
    "ax3 = fig.add_subplot(313)\n",
    "ax3.imshow(np.asarray(im))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Run data analytics <a class=\"anchor\" id=\"head-6\"></a>\n",
    "\n",
    "In this section, we run data analytics with the KITTI data. This service supports the following tasks:\n",
    "\n",
    "- analyze - This task analyzes the input files and generate graphs for calculated statistics. It can also generate the images with bounding boxes. The graphs can be generated locally or on wandb depend on the user input.\n",
    "\n",
    "- validate - This task validates the input files by calculating the invalid coordinates, imbalance data and suggests whether data needs to be revised."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.a Provide data analyze specification \n",
    "We provide specification file to configure the data analyze parameters including:\n",
    "\n",
    "- results_dir: \"/results/analytics\"\n",
    "- data:\n",
    "    - input_format: \"KITTI\"\n",
    "    - image_dir: \"/data/images\"\n",
    "    - ann_path: \"/data/labels\"\n",
    "- workers: 2\n",
    "- image:\n",
    "    - generate_image_with_bounding_box: False\n",
    "    - sample_size: 100\n",
    "- graph:\n",
    "    - generate_summary_and_graph: True\n",
    "    - height: 15\n",
    "    - width: 15\n",
    "    - show_all: False\n",
    "- wandb:\n",
    "    - visualize: False\n",
    "    - project: \"tao data analytics\"\n",
    " \n",
    "  \n",
    "- Image section configures if we want to generate images with bounding boxes.\n",
    "- Graph section configures the height and width of generated graphs. graph.show_all parameter decides if we want to visualize all the annotation data on generated graphs.\n",
    "- Wandb section configures if we want to generate the graph on wandb. By default all the graphs and images will be generated locally inside results_dir folder.\n",
    "\n",
    "Please refer to the TAO documentation about Data Analytics to get all the parameters that are configurable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/analytics.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.b Analyze data with Local Data visualization\n",
    "In below cell the annotation files will be analyzed and the insight graphs will be generated locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao dataset analytics analyze -e $SPECS_DIR/analytics.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Generated graphs:')\n",
    "print('---------------------')\n",
    "!ls -l $HOST_RESULTS_DIR/analytics/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below cell generates graphs as well as the images with bounding boxes locally by setting image.generate_image_with_bounding_box=True."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao dataset analytics analyze -e $SPECS_DIR/analytics.yaml image.generate_image_with_bounding_box=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Generated images with bounding boxes:')\n",
    "print('---------------------')\n",
    "!ls -l $HOST_RESULTS_DIR/analytics/image_with_bounding_boxes"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets visualize an image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image, ImageDraw\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "img = Image.open(os.path.join(os.environ[\"HOST_RESULTS_DIR\"], \"analytics/image_with_bounding_boxes/000015.png\"))\n",
    "\n",
    "# plot \n",
    "fig = plt.figure(figsize = (10,10))\n",
    "ax1 = fig.add_subplot(311)\n",
    "ax1.imshow(np.asarray(img))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.c Analyze data with Wandb visualization\n",
    "Here we are using wandb to visualize the data and uncover insights. Please refer the wandb integration documentation to setup wandb __[TAO Toolkit WandB Integration](https://tlt.gitlab-master-pages.nvidia.com/tlt-docs/text/mlops/wandb.html)__"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Steps to generate wandb login key.\n",
    "1. Create and login into wandb account.\n",
    "2. Go to user setting and copy the API Key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Please define wandb api key to enable wandb login.\n",
    "%env WANDB_API_KEY=WANDB_LOGIN_API_KEY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set docker env variable.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs[\"Envs\"]= [\n",
    "        {\n",
    "            \"variable\": \"WANDB_API_KEY\",\n",
    "            \"value\": os.environ[\"WANDB_API_KEY\"]\n",
    "        }\n",
    "    ]\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In below cell the annotation files will be analyzed and the insight graphs will be generated on wandb by setting wandb.visualize=True. Generated graphs can be find in wandb under the project name specified by wandb.project parameter.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "!tao dataset analytics analyze -e $SPECS_DIR/analytics.yaml wandb.visualize=True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below cell generates graphs as well as the annotated images with bounding boxes in wandb by setting image.generate_image_with_bounding_box=True and wandb.visualize=True."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao dataset analytics analyze -e $SPECS_DIR/analytics.yaml wandb.visualize=True image.generate_image_with_bounding_box=True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.d Provide data validate specification \n",
    "We provide specification file to configure the data validate parameters including:\n",
    "- apply_correction: True\n",
    "- results_dir: \"/results/analytics\"\n",
    "- data:\n",
    "    - input_format: \"KITTI\"\n",
    "    - image_dir: \"/data/images\"\n",
    "    - ann_path: \"/data/labels\"\n",
    "- workers: 2\n",
    "\n",
    "apply_correction parameter decides if we want to correct the invalid bounding box coordinates in the annotation files. The corrected files will be saved under the results_dir.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/validate.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6.e Data Validation\n",
    "Data validation tasks calculate the number of inverted and out of bound bounding boxes coordinates and suggests if user should proceed with training with given data or apply corrections. The validation summary prints the count of object tags also to review if the data is imbalanced. Invalid coordinates can be corrected by providing apply_correction = True. The corrected files will be saved in results_dir. Below are the correction conditions for bounding box coordinates.\n",
    "\n",
    "- Set negative coordinates to 0.\n",
    "\n",
    "- Swap the inverted coordinates.\n",
    "\n",
    "- If xmax is greater than image_width, then set xmax = image_width.\n",
    "\n",
    "- If ymax is greater than image_height, then set ymax = image_height.\n",
    "\n",
    "In below cell , the resultant number of invalid coordinates are 0 because there are no inverted or out of bound coordinates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao dataset analytics validate -e $SPECS_DIR/validate.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "5b3ded1ccb95c1d9bd405e7b823d9e85424cde40fbb5985eb47e999ef50e15b4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
